/**
 * Parallel multiplication of two square matrix using MPI library
 * @file matrix_multiply_mpi.c
 * @author Gvozdarev Sergey
 * Copyright 2012 graphite-mpi team
 * version 2
 */
#include <unistd.h>
#include <stdlib.h>
#include <stdio.h>

#include "mpi.h"

#define __DEBUG
#define MAT_SIZE 1000

#define G_ELEM_NUM 1
#define G_FINISH   -1
#define G_MAIN_RANK 0

#define  ijToN(i,j) ((i-1)*MAT_SIZE+j-1)
#define  NToi(N) (N/MAT_SIZE+1)
#define  NToj(N) (N%MAT_SIZE+1)

/// print matrix on screen
void printMat( double mat[]) 
{
    if ( MAT_SIZE > 15)
    {
        printf("Matrix is too large\n");
        return;
    }
    int i,j;
    for ( i = 1; i <= MAT_SIZE; i++ )
    {
        printf("||   ");
        for ( j = 1; j <= MAT_SIZE; j++)
        {
            printf("%.2f\t",mat[ ijToN(i,j)] );
        }
        printf("||\n");
    }
}
/// generating new random matrix
void genMat( double mat[] )
{
    int i;
    for ( i = 0; i < MAT_SIZE * MAT_SIZE; i++)
    {
        mat[ i] =   1.0 * ( (double)random() / (double)RAND_MAX );
    }
}
/// generating the identity matrix
void iMat( double mat[])
{
    int i;
    for ( i = 1; i <= MAT_SIZE; i++)
    {
        mat[ ijToN(i,i)] = 1.0;
    }
}

main( int argc, char ** argv )  
{
    double ttotal_begin, ttotal_end;
    double tinit_begin, tinit_end;
    double tbcast_begin, tbcast_end;
    double tcompute_begin, tcompute_end;


    double *mat1 = (double*)malloc( sizeof(double)*MAT_SIZE*MAT_SIZE);
    double *mat2 = (double*)malloc( sizeof(double)*MAT_SIZE*MAT_SIZE);
    double *mat3 = (double*)malloc( sizeof(double)*MAT_SIZE*MAT_SIZE);
    double *row  = (double*)malloc( sizeof(double)*MAT_SIZE);


    
    int comm_rank,comm_size,num_of_workers;
    MPI_Status status;

    //// Init
    MPI_Init( &argc, &argv ); 
    ttotal_begin = tinit_begin = MPI_Wtime();

    MPI_Comm_rank( MPI_COMM_WORLD, &comm_rank );
    MPI_Comm_size( MPI_COMM_WORLD, &comm_size );
    
    if ( comm_size == 1)
    {
        printf("comm_size must be bigger than 1");
        MPI_Finalize();
        return 0;
    }

    if ( comm_size > MAT_SIZE)
    {
        num_of_workers = MAT_SIZE;
    }
    else
    {
        num_of_workers = comm_size-1;
    }
    
    if ( comm_rank == 0) // main process
    {
        srandom( getpid());

        mat1 = (double*)malloc( sizeof(double)*MAT_SIZE*MAT_SIZE);
        mat2 = (double*)malloc( sizeof(double)*MAT_SIZE*MAT_SIZE);
        mat3 = (double*)malloc( sizeof(double)*MAT_SIZE*MAT_SIZE);
        row  = (double*)malloc( sizeof(double)*MAT_SIZE);

        if ( !mat1 || !mat2 || !mat3 || !row)
        {
            printf("Can't allocate memory\n");
            MPI_Finalize();
            return 0;
        }
        genMat( mat1);
        iMat( mat2);
        
#ifdef __DEBUG
        printf("Matrix1\n");
        printMat( mat1);
        printf("Matrix2\n");
        printMat( mat2);
#endif

        tinit_end = tbcast_begin = MPI_Wtime();
        MPI_Bcast( mat1, MAT_SIZE*MAT_SIZE, MPI_DOUBLE, G_MAIN_RANK,MPI_COMM_WORLD );
        MPI_Bcast( mat2, MAT_SIZE*MAT_SIZE, MPI_DOUBLE, G_MAIN_RANK,MPI_COMM_WORLD );
        tbcast_end = tcompute_begin = MPI_Wtime();
    
        
        int current  = 1;
        int progress = 1;
        int i;

        for ( i = 1; i <= num_of_workers; i++)
        {
            MPI_Send( &current, 1, MPI_INT, i, G_ELEM_NUM, MPI_COMM_WORLD);// send num of row to compute
            current++;
        }
        do
        {
            MPI_Recv( row, MAT_SIZE, MPI_DOUBLE, MPI_ANY_SOURCE, MPI_ANY_TAG,  MPI_COMM_WORLD,&status); // recieve result
            progress++;

            for ( i = 1; i <= MAT_SIZE; i++)
            {
                mat3[ ijToN( status.MPI_TAG, i)] = row[ i-1];
            }
            if( current <= MAT_SIZE)
            {
                MPI_Send( &current,1, MPI_INT, status.MPI_SOURCE , G_ELEM_NUM,MPI_COMM_WORLD );// send new task 
                current++;
            }
        }while ( progress <= MAT_SIZE);
        
        for ( i = 1; i < comm_size; i++)
        {
            int n = G_FINISH;
            MPI_Send( &n, 1, MPI_INT, i, G_ELEM_NUM, MPI_COMM_WORLD);
        }
        
    }else 
    { // workers

        mat1 = (double*)malloc( sizeof(double)*MAT_SIZE*MAT_SIZE);
        mat2 = (double*)malloc( sizeof(double)*MAT_SIZE*MAT_SIZE);
        row  = (double*)malloc( sizeof(double)*MAT_SIZE);
        if ( !mat1 || !mat2 || !row)
        {
            MPI_Finalize();
            return 0;
        }    
        MPI_Bcast( mat1, MAT_SIZE*MAT_SIZE, MPI_DOUBLE, G_MAIN_RANK,MPI_COMM_WORLD );
        MPI_Bcast( mat2, MAT_SIZE*MAT_SIZE, MPI_DOUBLE, G_MAIN_RANK,MPI_COMM_WORLD );
        
        int n;
        do
        {
            MPI_Recv( &n, 1, MPI_INT, G_MAIN_RANK, G_ELEM_NUM, MPI_COMM_WORLD, &status); // recieve num of raw
            if ( n == G_FINISH)
            {
                free( mat1);
                free( mat2);
                free( row);
                MPI_Barrier( MPI_COMM_WORLD);
                MPI_Finalize();
                return 0;
            }

            int c;
            for ( c = 1; c<=MAT_SIZE; c++)
            {
                int k;
                int i = n;
                int j = c;
                row[ c-1] = 0;
                for ( k = 1; k <= MAT_SIZE; k++ )
                {
                    row[ c-1] += mat1[ ijToN( i,k)] * mat2[ ijToN( k,j)];
                }
            }
            MPI_Send( row, MAT_SIZE ,MPI_DOUBLE, G_MAIN_RANK, n, MPI_COMM_WORLD);
        }
        while(1);
    }
#ifdef __DEBUG
        printf("Resutlt\n");
        printMat( mat3);
#endif

    ttotal_end = tcompute_end = MPI_Wtime();

    printf( "--------------------------------\n");
    printf("Matrix size %d * %d\n",MAT_SIZE,MAT_SIZE);
    printf("Number of workers %d\n",num_of_workers);
    printf( "--------------------------------\n");
    printf( "|| Init time     \t%.2f ms \n", ( tinit_end  - tinit_begin)*1000);
    printf( "|| Broadcast time\t%.2f ms \n", ( tbcast_end - tbcast_begin)*1000);
    printf( "|| Compute   time\t%.2f ms \n", ( tcompute_end - tcompute_begin)*1000);
    printf( "--------------------------------\n");
    printf( "|| Total time    \t%.2f ms \n", ( ttotal_end - ttotal_begin)*1000);
    printf( "--------------------------------\n");

    free( mat1);
    free( mat2);
    free( mat3);
    free( row );
    
    MPI_Barrier( MPI_COMM_WORLD);
    MPI_Finalize();
} 